package com.moran.config;

import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

/**
 * 向量模型配置类
 *
 */
@Configuration
public class EmbeddingConfig {

    @Value("${spring.ai.zhipuai.base-url}")
    private String baseUrl;

    @Value("${spring.ai.zhipuai.api-key}")
    private String token;

    @Value("${spring.ai.zhipuai.embedding.options.model}")
    private String model;

    /**
     * 实例化智普AI向量模型对象
     *
     * @return
     */
    @Bean
    @Primary
    public EmbeddingModel embeddingModel() {
        ZhiPuAiApi zhiPuAiApi = new ZhiPuAiApi(baseUrl, token);

        ZhiPuAiEmbeddingOptions options = ZhiPuAiEmbeddingOptions.builder()
                .model(model)
                .build();

        ZhiPuAiEmbeddingModel embeddingModel = new ZhiPuAiEmbeddingModel(zhiPuAiApi, MetadataMode.EMBED, options);
        return embeddingModel;
    }

    /**
     * 实例化简单的内存向量存储对象
     *
     * @param embeddingModel
     * @return
     */
    @Bean
    @Primary
    public SimpleVectorStore vectorStore(EmbeddingModel embeddingModel) {
        // 创建向量存储
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        return vectorStore;
    }

}
